"""
单文件夹推理脚本 - 针对3t3数据集优化
基于 infer_on_folder.py 修改
"""

import os
import torch
import tifffile
import numpy as np
import json
from tqdm import tqdm

from embedtrack.infer.inference import init_model
from embedtrack.utils.clustering import Cluster

# ========== 用户需填写以下路径 ==========
# 模型设置
MODEL_NAME = "adam_norm_onecycle_24"  # 您的模型名称
DATA_SET = "3t3"  # 数据集名称

# 自动获取最新的模型文件
MODEL_BASE_PATH = f'models/{DATA_SET}/{MODEL_NAME}'
if os.path.exists(MODEL_BASE_PATH):
    # 获取最新的时间戳目录
    import datetime
    timestamps = [d for d in os.listdir(MODEL_BASE_PATH) 
                  if os.path.isdir(os.path.join(MODEL_BASE_PATH, d))]
    timestamps.sort()
    latest_timestamp = timestamps[-1]
    
    MODEL_PATH = os.path.join(MODEL_BASE_PATH, latest_timestamp, 'best_iou_model.pth')
    CONFIG_PATH = os.path.join(MODEL_BASE_PATH, latest_timestamp, 'config.json')
    
    print(f"使用模型: {MODEL_PATH}")
    print(f"配置文件: {CONFIG_PATH}")
else:
    print(f"错误: 模型目录不存在 {MODEL_BASE_PATH}")
    exit(1)

# 输入输出路径设置
INPUT_FOLDER = 'ctc_raw_data/train/3t3/46/images'  # 输入图片文件夹
OUTPUT_FOLDER = 'results/3t3/46/masks'  # 输出mask文件夹

# 检查输入路径
if not os.path.exists(INPUT_FOLDER):
    print(f"错误: 输入文件夹不存在 {INPUT_FOLDER}")
    print("请检查路径或修改 INPUT_FOLDER 变量")
    exit(1)

# 检查模型文件
if not os.path.exists(MODEL_PATH):
    print(f"错误: 模型文件不存在 {MODEL_PATH}")
    exit(1)

if not os.path.exists(CONFIG_PATH):
    print(f"错误: 配置文件不存在 {CONFIG_PATH}")
    exit(1)

print(f"输入文件夹: {INPUT_FOLDER}")
print(f"输出文件夹: {OUTPUT_FOLDER}")

# ========== 加载配置 ==========
with open(CONFIG_PATH, 'r') as f:
    config = json.load(f)

model_class = config['model_dict']['name']
input_channels = config['model_dict']['kwargs']['input_channels']
n_classes = config['model_dict']['kwargs']['n_classes']
grid_y = config['grid_dict']['grid_y']
grid_x = config['grid_dict']['grid_x']
pixel_y = config['grid_dict']['pixel_y']
pixel_x = config['grid_dict']['pixel_x']
crop_size = config['train_dict']['crop_size']
min_mask_size = config['train_dict'].get('min_mask_size', 10)

print(f"模型配置:")
print(f"  模型类型: {model_class}")
print(f"  输入通道: {input_channels}")
print(f"  类别数: {n_classes}")
print(f"  网格大小: {grid_y}x{grid_x}")
print(f"  像素大小: {pixel_y}x{pixel_x}")
print(f"  裁剪大小: {crop_size}")
print(f"  最小掩码大小: {min_mask_size}")

# ========== 初始化模型 ==========
model_dict = {
    'kwargs': {
        'input_channels': input_channels,
        'n_classes': n_classes,
    }
}
model_configs = {
    'model_class': model_class,
    'model_cktp_path': MODEL_PATH
}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

try:
    model = init_model(model_dict, model_configs)
    model = model.to(device)
    model.eval()
    print("✅ 模型加载成功")
except Exception as e:
    print(f"❌ 模型加载失败: {e}")
    exit(1)

# ========== 初始化Cluster ==========
try:
    cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x)
    cluster = cluster.to(device)
    print("✅ Cluster初始化成功")
except Exception as e:
    print(f"❌ Cluster初始化失败: {e}")
    exit(1)

# ========== 创建输出文件夹 ==========
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
print(f"输出文件夹已创建: {OUTPUT_FOLDER}")

# ========== 推理每张图片 ==========
img_list = [f for f in os.listdir(INPUT_FOLDER) 
            if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg', '.jpeg'))]

if len(img_list) == 0:
    print(f"警告: 在 {INPUT_FOLDER} 中没有找到图片文件")
    exit(1)

print(f"找到 {len(img_list)} 张图片，开始推理...")

success_count = 0
error_count = 0

for img_name in tqdm(img_list, desc='推理中'):
    try:
        img_path = os.path.join(INPUT_FOLDER, img_name)
        image = tifffile.imread(img_path)
        
        # 处理图像维度
        if image.ndim == 2:
            image = image[np.newaxis, ...]  # (1, H, W)
        elif image.ndim == 3 and image.shape[0] != input_channels:
            image = np.transpose(image, (2, 0, 1))  # (C, H, W)
        
        input_tensor = torch.from_numpy(image).float().unsqueeze(0).to(device)  # (1, C, H, W)
        
        with torch.no_grad():
            pred = model(input_tensor)
            if isinstance(pred, (tuple, list)):
                pred = pred[0]  # 取分割输出
            mask = cluster.cluster_pixels(pred.squeeze(0), n_sigma=2, min_obj_size=min_mask_size)
        
        mask = mask.cpu().numpy().astype('uint16')
        output_path = os.path.join(OUTPUT_FOLDER, os.path.splitext(img_name)[0] + '_mask.tif')
        tifffile.imwrite(output_path, mask)
        
        success_count += 1
        
    except Exception as e:
        print(f"\n❌ 处理图片 {img_name} 时出错: {e}")
        error_count += 1
        continue

print(f"\n{'='*50}")
print("推理完成!")
print(f"成功处理: {success_count} 张图片")
print(f"处理失败: {error_count} 张图片")
print(f"结果已保存到: {OUTPUT_FOLDER}")
print(f"{'='*50}") 